import java.math.BigInteger;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.security.spec.RSAPrivateKeySpec;
import java.util.ArrayList;
import java.util.Date;
/**
 * Secure server (provides mutual authentication and AES block encryption using Diffie Hellman key exchange) 
 * @author benjamin
 */
public class Server implements SocketListener,Runnable{
	private Thread thread;
	private TCPServerSocket sock;
	private NonBlockingReader rdr;
	private ArrayList<TCPSocket> unknownSockets;
	private ArrayList<TCPSocket> socketWithNoExchangedKeyArray;
	private ArrayList<Pair<TCPSocket,String>>authenticatedSockets; 
	private ArrayList<Pair<TCPSocket,Pair<String,byte[]>>> unauthenticatedSockets;
	private KeyringReader keyring;
	private CypherMachine cypher;
	private RSAPrivateKeySpec pk;
	/**
	 * Default constructor for server program
	 */
	public Server(){
		//Try to open a listening socket on the system
		try {
			sock = new TCPServerSocket(this,ProtocolInfo.SERVER_PORT);
		} catch (Exception e) {
			System.out.println("Could not establish socket. This normally happens when you try to run multiple servers.");
			e.printStackTrace();
			System.exit(1);
		}
		//Socket opened, instantiate a non-blocking input reader
		rdr = new NonBlockingReader(Driver.s);
		//Read the keyring
		try{
			keyring = new KeyringReader(ProtocolInfo.KEYRING_LOCATION);
		} catch (Exception e) {
			System.out.println("\nCould not read keyring. Missing file or corrupted");
			e.printStackTrace();
			System.exit(1);
		}
		//Keyring loaded, now instantiate our cypher machine
		try{
			cypher = new CypherMachine();
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		}
		//Here is the RSA private key of the server (this should probably not be stored in memory like this.... but ye its fine for our purposes...
		pk = new RSAPrivateKeySpec(
				new BigInteger("106552646548987747775799982470752321354613464888190003586735280912689621477051251518731532849179225407091383299774564288466375737201985922415443230079706656965365814775533860604483929010866271298722865213854049246766680001231445963555353946101665118756527053733522791467610722357913068363694616394809975353403"), 
				new BigInteger("106227478639107870003623044306140720819481818749708870780574683751047527327251745128391018230416757715841886425337911514343463868373074053021313369908872735808120018916948511985165473414468692763600577631646191044252687002911039909747897627609423694632590367687176641827882267809786761047246539003414789701601"));
		//fire up the server main loop
		thread = new Thread(this);
		thread.start();
		unknownSockets = new ArrayList<TCPSocket>();
		socketWithNoExchangedKeyArray = new ArrayList<TCPSocket>();
		unauthenticatedSockets = new ArrayList<Pair<TCPSocket,Pair<String,byte[]>>>(); 
		authenticatedSockets = new ArrayList<Pair<TCPSocket,String>>();
		System.out.println("\n***********************************************");
		System.out.println("*                   Server                    *");
		System.out.println("***********************************************");
		prompt();
		//Wait for the server thread to die before joining up with the main thread
		try {
			thread.join();
		} catch (InterruptedException e) {}
		sock.interrupt();
		rdr.interrupt();
	}
	@Override
	public synchronized void onIncommingData(TCPSocket sock, String data) {
		//If the client is an unknown socket then decrypt timestamp(C) using server private key and encrypt timestamp(C),timestamp(S) using client public key  
		for (int i = 0; i < unknownSockets.size(); ++i){
			TCPSocket s = unknownSockets.get(i);
			if (s == sock){
				try{
					String name = data.substring(0,data.indexOf(','));
					byte[] timestamp = cypher.RSAPriKeyDecrypt(
							Base64.decode(data.substring(data.indexOf(',')+1).getBytes()),pk);
					byte[] serverStamp = ByteBuffer.allocate(8).putLong(new Date().getTime()).array();
					System.out.println("\nAUTH: " + name + " has sent me a time stamp");
					System.out.println("\nAUTH: I'm authenticating by sending it back. I'm also sending along my own time stamp for this client");
					Pair<TCPSocket,Pair<String,byte[]>> pair = new Pair<TCPSocket,Pair<String,byte[]>>(s,
							new Pair<String,byte[]>(name,serverStamp));
					try {
						s.sendData(Base64.encodeBytes(cypher.RSAPubKeyEncrypt(timestamp,
								keyring.getKeys().get(name)))+','+Base64.encodeBytes(cypher.RSAPubKeyEncrypt(serverStamp,
										keyring.getKeys().get(name))));
						unauthenticatedSockets.add(pair);
						return;
					} catch (Exception e){
						System.out.println("\nLost connection to " + name + ". The party will have to retry later.");
					}
				} catch (Exception e){
					System.out.println("AUTH: Notion of mistrust raised for "+sock.clientSocket.getInetAddress().toString());
					try {s.sendData(ProtocolInfo.NO_TRUST); } catch (Exception exp) {};
				}
				finally {
					unknownSockets.remove(i--);
					prompt();
				}
			}
		}
		//incomming timestamp(S) from client decrypt using public key of client and check for a match. Send ACK if this is the case
		for (int i = 0; i < unauthenticatedSockets.size(); ++i){
			Pair<TCPSocket,Pair<String,byte[]>> s = unauthenticatedSockets.get(i);
			if (s.getVal1() == sock){
				try{
				byte[] decryption = cypher.RSAPubKeyDecrypt(Base64.decode(data), keyring.getKeys().get(s.getVal2().getVal1()));
				if (CypherMachine.compareByteArrays(decryption, s.getVal2().getVal2())){
					try {
						s.getVal1().sendData(ProtocolInfo.HANDSHAKE_ACK);
						System.out.println("\nAUTH: " + s.getVal2().getVal1() + " has authenticed. Handshake complete.");
						authenticatedSockets.add(new Pair<TCPSocket, String>(s.getVal1(), s.getVal2().getVal1()));
						return;
					} catch (Exception e){
						System.out.println("\nLost connection to " + s.getVal2().getVal1() + ". The party will have to retry later.");
						unauthenticatedSockets.remove(i--);
					}
					finally{
						prompt();
					}
				} else throw new Exception("I do not trust this connection");
				} catch (Exception e) {
					System.out.println("\nAUTH: Notion of mistrust raised for "+sock.clientSocket.getInetAddress().toString());
					e.printStackTrace();
					prompt();
					//try {s.getVal1().sendData(ProtocolInfo.NO_TRUST); } catch (Exception exp) {};
				}
				finally{
					unauthenticatedSockets.remove(i--);
				}
			}
		}
		//Authenticated sockets can just chat with no probs :)
		for (int i = 0; i < authenticatedSockets.size(); ++i){
			Pair<TCPSocket,String> s = authenticatedSockets.get(i);
			if (s.getVal1() == sock){
				System.out.println("\n"+s.getVal2()+" says: "+data);
				prompt();
				return;
			}
		}
		//ignore any other message
	}
	/**
	 * Handles prompt printouts
	 */
	private void prompt(){
		System.out.println("\nType 'X' to exit");
		System.out.print(">");
	}
	@Override
	public void run() {
		while (!Thread.interrupted()){
			String input = rdr.getNextLine();
			if (input != null){
				if (input.equals("X"))
					thread.interrupt();
				else{
					System.out.println("\nInvalid Input");
					prompt();
				}
			}
			try{
				Thread.sleep(10);
			}
			catch (InterruptedException e){break;}
		
		}
	}
	@Override
	public synchronized void onIncommingConnection(TCPSocket s) {
		for (int i = 0; i < unknownSockets.size(); ++i)
			if (unknownSockets.get(i).clientSocket.getInetAddress().equals(s.clientSocket.getInetAddress()))
				unknownSockets.remove(i--);
		unknownSockets.add(s);
	}
	@Override
	public void onClientSecured(TCPSocket s) {
		prompt();
	}
	@Override
	public void onClientDisconnect(TCPSocket s) {
		//remove from unknown sockets
		if (this.unknownSockets.contains(s))
			this.unknownSockets.remove(s);
		//remove from unauthenticated sockets
		Pair<TCPSocket,Pair<String,byte[]>> removable = null;
		for (Pair<TCPSocket,Pair<String,byte[]>> p:this.unauthenticatedSockets){
			if (p.getVal1() == s)
				removable = p;
		}
		if (removable != null)
			this.unauthenticatedSockets.remove(removable);
		//remove from authenticated sockets
		Pair<TCPSocket,String> removableAut = null;
		for (Pair<TCPSocket,String> p:this.authenticatedSockets){
			if (p.getVal1() == s)
				removableAut = p;
		}
		if (removableAut != null)
			this.authenticatedSockets.remove(removableAut);
	}
}
